package com.briup.facade.service.impl;

import com.briup.common.properties.ZookeeperProperties;
import com.briup.common.web.SecurityUtil;
import com.briup.dataset.mapper.extend.DataSetExtendMapper;
import com.briup.facade.service.IFascadeTrainingService;
import com.briup.model.entity.Model;
import com.briup.model.entity.Training;
import com.briup.model.entity.TrainingLabel;
import com.briup.model.entity.extend.TrainingVo;
import com.briup.model.entity.extend.TraningEx;
import com.briup.model.mapper.ModelMapper;
import com.briup.model.mapper.TrainingLabelMapper;
import com.briup.model.mapper.TrainingMapper;
import com.briup.model.mapper.extend.ModelExtendMapper;
import com.briup.model.mapper.extend.TrainingExtendMapper;
import com.briup.model.vo.ZookeeperInitVo;
import com.briup.model.vo.ZookeeperOptimizeVo;
import com.briup.operator.entity.Operator;
import com.briup.operator.service.IOpeatorService;
import com.briup.upload.service.IRecoverImageService;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.curator.framework.CuratorFramework;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.*;

/**
 * facade的训练服务
 */
@Service
public class FascadeTrainingServiceImpl implements IFascadeTrainingService {
    @Autowired
    private TrainingMapper trainingMapper;
    @Autowired
    private TrainingLabelMapper trainingLabelMapper;
    @Autowired
    private TrainingExtendMapper trainingExtendMapper;
    @Autowired
    private IRecoverImageService recoverImageService;
    @Autowired
    private IOpeatorService opeatorService;
    @Autowired
    private DataSetExtendMapper dataSetExtendMapper;
    @Autowired
    private ModelMapper modelMapper;
    @Autowired
    private ModelExtendMapper modelExtendMapper;
    // 注入zookeeper
    @Autowired
    private CuratorFramework client;
    @Autowired
    private ZookeeperProperties zookeeperProperties;

    @Transactional
    @Override
    public Integer insertTraining(Training training) {
        trainingMapper.insertSelective(training);
        // 返回主键回填
        return training.getId();
    }
    @Transactional
    @Override
    public void insertTrainingLabel(Integer trainingId, List<TrainingLabel> trainingLabelList) {
        trainingLabelList.forEach(trainingLabel -> {
            trainingLabel.setTrainingId(trainingId);
            trainingLabelMapper.insertSelective(trainingLabel);
        });
    }
    @Transactional
    @Override
    public void insertTrainingDataset(Integer trainingId, List<Integer> datasetIdList) {
        datasetIdList.forEach(datasetId->{
            trainingExtendMapper.insertTrainingAndDatasets(trainingId,datasetId);
        });
    }

    @Transactional
    @Override
    public void trainingModel(TrainingVo trainingVo, Integer modelId) {
        // 循环遍历前构建分类名去重集合
        Set<String> labelNameSet = new HashSet<>();
        // 操作数据集
        trainingVo.getDatasetIdList().forEach(dataSetId -> {
            // 先相同对多数据集中分类名称去重(取消)
            distinctLabel(dataSetId, labelNameSet);
            // 顺便还原数据集到指定位置,方便训练使用
            // recoverImageService.recoverImage(dataSetId);
        });

        // --------- 2.封装zookeeper需要的数据(开始) -----------------
        ObjectMapper objectMapper = new ObjectMapper();
        if (trainingVo.getTaskType() == 0) {
            ZookeeperInitVo zookeeperInitVo = fillZookeeperInitVo(modelId,trainingVo,labelNameSet);
            String zookeeperInitVoJson = null;
            try {
                zookeeperInitVoJson = objectMapper.writeValueAsString(zookeeperInitVo);
                client.setData().forPath(zookeeperProperties.getTaskPath(), zookeeperInitVoJson.getBytes());
            } catch (Exception e) {
                e.printStackTrace();
            }
        } else if (trainingVo.getTaskType() == 1) {
            // 优化训练需要的参数
            // 通过模型id和最后一个模型版本定位到需要优化的训练记录信息
            ZookeeperOptimizeVo zookeeperOptimizeVo = fillZookeeperOptimizeVo(modelId,trainingVo,labelNameSet);
            try {
                String zookeeperOptimizerVoJson = objectMapper.writeValueAsString(zookeeperOptimizeVo);
                client.setData().forPath(zookeeperProperties.getTaskPath(), zookeeperOptimizerVoJson.getBytes());
            } catch (Exception e) {
                e.printStackTrace();
            }
            // 改变模型状态  3训练中
            Model model = modelMapper.selectByPrimaryKey(modelId);
            model.setTrainingStatus(3);
            model.setId(modelId);
            modelMapper.updateByPrimaryKeySelective(model);
        }
        // --------- 2.封装zookeeper需要的数据(完成) -----------------
    }

    private ZookeeperOptimizeVo fillZookeeperOptimizeVo(Integer modelId, TrainingVo trainingVo, Set<String> labelNameSet) {
        // 任务id
        String taskId = UUID.randomUUID().toString();
        // 用户id
        Integer userId = Integer.parseInt(SecurityUtil.getUserId());
        //  模型信息
        Model model = modelMapper.selectByPrimaryKey(modelId);
        String lastModelVersion = trainingVo.getLastModelVersion();
        TraningEx traningEx = modelExtendMapper.findTraningExByModelIdAndVersion(modelId, lastModelVersion);
        ZookeeperOptimizeVo zookeeperOptimizeVo = new ZookeeperOptimizeVo();
        zookeeperOptimizeVo.setLastVersion(lastModelVersion); // 最后一个模型版本
        zookeeperOptimizeVo.setOldModelPath(traningEx.getModelFileaddr()); // 原模型文件地址 从数据库中拿
        zookeeperOptimizeVo.setUserId(userId); // 当前用户id
        zookeeperOptimizeVo.setTaskId(taskId);  // 任务id 随机生成字符串
        zookeeperOptimizeVo.setTaskTime(new Date()); // 任务时间
        zookeeperOptimizeVo.setTaskType(trainingVo.getTaskType()); // 任务类型
        zookeeperOptimizeVo.setDatasetIdList(trainingVo.getDatasetIdList()); // 训练所用数据集id
        zookeeperOptimizeVo.setDatasetAddress("Intergrationfiles/" + taskId); // 训练所需数据集地址
        zookeeperOptimizeVo.setModelId(modelId); // 训练模型id
        zookeeperOptimizeVo.setModelName(model.getModelName()); // 训练模型名称
        zookeeperOptimizeVo.setNewModelPath("ModelFiles/" + userId + "/" + model.getModelName()); // 模型文件保存地址
        zookeeperOptimizeVo.setStatus(0);   // 训练状态,由训练程序控制,只传0
        // 训练需要入库的数据
        zookeeperOptimizeVo.setLabelNameSet(labelNameSet);
        return zookeeperOptimizeVo;
    }

    private ZookeeperInitVo fillZookeeperInitVo(int modelId,TrainingVo trainingVo,Set<String> labelNameSet) {
        // 初始化训练需要的参数
        // 任务id
        String taskId = UUID.randomUUID().toString();
        // 用户id
        Integer userId = Integer.parseInt(SecurityUtil.getUserId());
        //  模型信息
        Model model = modelMapper.selectByPrimaryKey(modelId);
        // 网络结构算子
        Operator newworkStructureInfo = opeatorService.findByUserIdAndOperatorUrl(userId, trainingVo.getNetWorkStructure());
        // 损失函数算子
        Operator lossValueInfo = opeatorService.findByUserIdAndOperatorUrl(userId, trainingVo.getLossValue());
        // 优化器算子
        Operator optimizerInfo = opeatorService.findByUserIdAndOperatorUrl(userId, trainingVo.getOptimizer());
        ZookeeperInitVo zookeeperInitVo = new ZookeeperInitVo();
        zookeeperInitVo.setLastVersion("V0"); // 初始化版本为V0
        zookeeperInitVo.setReslution(trainingVo.getReslution());  // 分辨率
        zookeeperInitVo.setIterateTimes(trainingVo.getIterateTimes()); // 迭代次数
        zookeeperInitVo.setOperatorType(newworkStructureInfo.getOperatorType()); // 算子类型 类型由网络结构决定 ???
        zookeeperInitVo.setNetworkStructure(newworkStructureInfo.getOperatorUrl()); // 网络结构
        zookeeperInitVo.setOptimizer(optimizerInfo.getOperatorUrl());  // 优化器
        zookeeperInitVo.setLossValue(lossValueInfo.getOperatorUrl()); // 损失函数
        zookeeperInitVo.setUserId(userId); // 当前用户id
        zookeeperInitVo.setTaskId(taskId); // 随机数,唯一标识任务
        zookeeperInitVo.setTaskTime(new Date()); // 任务时间
        zookeeperInitVo.setTaskType(trainingVo.getTaskType()); // 任务类型  0 初始化训练 1 优化训练
        zookeeperInitVo.setDatasetIdList(trainingVo.getDatasetIdList());  // 训练所用数据集id
        zookeeperInitVo.setDatasetAddress("Intergrationfiles/" + taskId);   // 训练所需数据集地址
        zookeeperInitVo.setModelId(modelId); // 训练模型id
        zookeeperInitVo.setModelName(model.getModelName()); // 训练模型名称
        zookeeperInitVo.setNewModelPath("ModelFiles/" + userId + "/" + model.getModelName()); // 模型文件保存地址
        zookeeperInitVo.setStatus(0); // 训练状态,由训练程序控制,只传0
        // 携带需要入库的数据给mock
        zookeeperInitVo.setLabelNameSet(labelNameSet);
        return zookeeperInitVo;
    }

    @Override
    public void cancelTraining(Integer modelId) {
        // 根据taskId删除zookeeper中的节点
        try {
            client.setData().forPath(zookeeperProperties.getTaskPath(), null);
            // 修改模型状态
            Model model = new Model();
            model.setId(modelId);
            model.setTrainingStatus(5); // 训练失败
            modelMapper.updateByPrimaryKeySelective(model);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    // 给分类名称去重
    private void distinctLabel(Integer dataSetId, Set<String> labelNameSet) {
        // 获取该数据集下所有的分类名称
        List<String> classNameList = dataSetExtendMapper.selectClassifyNameByDatasetId(dataSetId);
        classNameList.forEach(className -> {
            labelNameSet.add(className);
        });
    }
}
